import random
import math

class GenerateNetwork:
    def get_unconnect_pair(edges, numNodes):
        map_arr = [0] * (numNodes + 1)
        node_queue_arr = [0] * (numNodes + 1)
        start_pos = 0
        end_pos = 0
        node_num = 0
        startBfsNode = edges[0][0]
        node_queue_arr[end_pos] = startBfsNode
        end_pos += 1
        node_num = 1
        map_arr[edges[0][0]] = 1

        while start_pos < end_pos:
            tmp = node_queue_arr[start_pos]
            start_pos += 1
            for edge in edges:
                startID = edge[0]
                endID = edge[1]
                if startID == tmp and endID != tmp and map_arr[endID] == 0:
                    node_queue_arr[end_pos] = endID
                    end_pos += 1
                    node_num += 1
                    map_arr[endID] = 1
                elif endID == tmp and startID != tmp and map_arr[startID] == 0:
                    node_queue_arr[end_pos] = startID
                    end_pos += 1
                    node_num += 1
                    map_arr[startID] = 1

        if node_num == numNodes:
            return False, (0, 0)
        else:
            for nodeIndex in range(1, numNodes):
                if map_arr[nodeIndex] + map_arr[nodeIndex + 1] == 1:
                    return True, (nodeIndex, nodeIndex + 1)
            return False, (0, 0)

    def make_graph_connect(edges, numNodes):
        is_unconnected, unconnected_pair = GenerateNetwork.get_unconnect_pair(edges, numNodes)
        while is_unconnected:
            edges.append(unconnected_pair)
            is_unconnected, unconnected_pair = GenerateNetwork.get_unconnect_pair(edges, numNodes)
    
    def generate_network_by_nodes_edges(
        num_nodes = 300, 
        num_edges = 800, 
        edge_capacity_min = 20, 
        edge_capacity_max = 100):
        outfile = open("generate.txt", "w")
        def distance(p1, p2):
            return math.sqrt((p1[0] - p2[0])**2 + (p1[1] - p2[1])**2)
        nodes = {i: (random.uniform(0, 1), random.uniform(0, 1)) for i in range(1, num_nodes + 1)}
        edges = []
        while len(edges) < num_edges:
            n1, n2 = random.sample(nodes.keys(), 2)
            d = distance(nodes[n1], nodes[n2])
            probability = math.exp(-d * 10)  
            if random.random() < probability:
                if n1 > n2: n1, n2 = n2, n1  
                edges.append((n1, n2))
        outfile.write(str(num_nodes) + "\n")
        for i in range(1, num_nodes + 1):
            outfile.write(str(i) + "\n")
        GenerateNetwork.make_graph_connect(edges, num_nodes)
        outfile.write(str(len(edges)) + "\n")
        for edge in edges:
            randomInt = random.randint(edge_capacity_min, edge_capacity_max)
            outfile.write(f"{edge[0]} {edge[1]} {randomInt}\n")

    def generate_network_with_var_capacity(numNodes, alpha, beta, run_round_index):
        outfile = open("generate.txt", "w")
        outfile.write(str(numNodes) + "\n")
        for i in range(1, numNodes+1):
            outfile.write(str(i) + "\n")
        generator = random.Random(1)
        edges = []
        for i in range(1, numNodes+1):
            for j in range(i + 1, numNodes+1):
                dist = math.sqrt((i - j) ** 2)
                prob = alpha * math.exp(-beta * dist)
                if generator.random() < prob:
                    edges.append((i, j))
        
        if len(edges) == 0:
            edges.append((1, 2))

        GenerateNetwork.make_graph_connect(edges, numNodes)
        outfile.write(str(len(edges)) + "\n")

        for tmp in edges:
            randomInt = random.randint(10 * (1 + run_round_index) - 10, 10 * (1 + run_round_index) + 10)
            outfile.write(f"{tmp[0]} {tmp[1]} {randomInt}\n")

        outfile.close()
        # print("数据已成功写入文件！")

    def generate_network(numNodes, alpha, beta, seed):
        outfile = open("generate.txt", "w")
        outfile.write(str(numNodes) + "\n")
        for i in range(1, numNodes+1):
            outfile.write(str(i) + "\n")
        generator = random.Random()
        edges = []
        for i in range(1, numNodes+1):
            for j in range(i + 1, numNodes+1):
                dist = math.sqrt((i - j) ** 2)
                prob = alpha * math.exp(-beta * dist)
                if generator.random() < prob:
                    edges.append((i, j))
        
        if len(edges) == 0:
            edges.append((1, 2))

        GenerateNetwork.make_graph_connect(edges, numNodes)
        outfile.write(str(len(edges)) + "\n")

        for tmp in edges:
            randomInt = random.randint(5, 100)
            outfile.write(f"{tmp[0]} {tmp[1]} {randomInt}\n")

        outfile.close()
        # print("数据已成功写入文件！")

